package com.qfedu.entity.chapter2.course09;

import java.io.IOException;
import java.util.Random;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.locks.ReentrantLock;


public class TokenBucket {
    // 平均流量(每秒生成的令牌数量）
    private int avgFlowRate = 512;
    // 一个令牌允许的数据包大小 1字节
    private int everyTokenSize = 1;
    // 桶内令牌上限1024个，最大流量峰值（瞬间最大流量）为 everyTokenSize*maxFlowRate =1024Byte=1K
    private int maxFlowRate = 1024;
    // 队列来缓存桶数量：
    private ArrayBlockingQueue<Byte> tokenQueue = new ArrayBlockingQueue<>(maxFlowRate);
    //任务调度，固定速率产生令牌
    private ScheduledExecutorService scheduledExecutorService = Executors
            .newSingleThreadScheduledExecutor();
    private volatile boolean isStart = false;
    private ReentrantLock lock = new ReentrantLock(true);

    public TokenBucket() {
    }

    public TokenBucket(int maxFlowRate, int avgFlowRate) {
        this.maxFlowRate = maxFlowRate;
        this.avgFlowRate = avgFlowRate;
    }

    public TokenBucket(int everyTokenSize, int maxFlowRate, int avgFlowRate) {
        this.everyTokenSize = everyTokenSize;
        this.maxFlowRate = maxFlowRate;
        this.avgFlowRate = avgFlowRate;
    }

    public static TokenBucket newBuilder() {
        return new TokenBucket();
    }

    /**
     * 获取足够的令牌个数
     *
     * @return
     */
    public boolean getTokens(byte[] dataSize) {
        int needTokenNum = dataSize.length / everyTokenSize + 1;// 传输内容大小对应的桶个数
        final ReentrantLock lock = this.lock;
        lock.lock();
        try {
            boolean result = needTokenNum <= tokenQueue.size(); // 是否存在足够的桶数量
            System.out.println("桶内令牌数" + (result ? "充足" : "不足") + "，总令牌数：" +
                    tokenQueue.size() + " 需要：" + needTokenNum);
            if (!result) {
                return false;
            }
            int tokenCount = 0;
            for (int i = 0; i < needTokenNum; i++) {
                Byte poll = tokenQueue.poll();
                if (poll != null) {
                    tokenCount++;
                }
            }
            return tokenCount == needTokenNum;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 添加令牌
     *
     * @param tokenNum
     */
    public void addTokens(Integer tokenNum) {
        for (int i = 0; i < tokenNum; i++) {
            // 若是桶已经满了，就不再加入新的令牌
            tokenQueue.offer((byte) 1);
        }
    }

    /**
     * build，完成初始化，并开始任务调度（始按照固定速率生产令牌）
     *
     * @return
     */
    public TokenBucket build() {
        start();
        return this;
    }

    public void start() {
        // 初始化桶队列大小
        if (maxFlowRate != 0) {
            tokenQueue = new ArrayBlockingQueue<Byte>(maxFlowRate);
        }
        // 初始化令牌生产者
        TokenProducer tokenProducer = new TokenProducer(avgFlowRate, this);
        scheduledExecutorService.scheduleAtFixedRate(tokenProducer, 0, 1,
                TimeUnit.SECONDS);
        isStart = true;
    }

    public void stop() {
        isStart = false;
        scheduledExecutorService.shutdown();
    }

    public boolean isStarted() {
        return isStart;
    }

    public TokenBucket everyTokenSize(int everyTokenSize) {
        this.everyTokenSize = everyTokenSize;
        return this;
    }

    public TokenBucket maxFlowRate(int maxFlowRate) {
        this.maxFlowRate = maxFlowRate;
        return this;
    }

    public TokenBucket avgFlowRate(int avgFlowRate) {
        this.avgFlowRate = avgFlowRate;
        return this;
    }

    /**
     * 定时任务，产生avgFlowRate个令牌
     */
    static class TokenProducer implements Runnable {
        private int avgFlowRate;
        private TokenBucket tokenBucket;

        TokenProducer(int avgFlowRate, TokenBucket tokenBucket) {
            this.avgFlowRate = avgFlowRate;
            this.tokenBucket = tokenBucket;
        }

        @Override
        public void run() {
            tokenBucket.addTokens(avgFlowRate);
        }
    }

    public static void main(String[] args) throws IOException, InterruptedException {
        tokenTest();
    }

    private static void tokenTest() throws InterruptedException, IOException {
        // 平均流量512，瞬时流量1024
        TokenBucket tokenBucket = TokenBucket.newBuilder().avgFlowRate(512)
                .maxFlowRate(1024).build();
        String data = "0000";// 四个字节
        for (int i = 1; i <= 20; i++) {
            //生成随机大小的流量请求
            int copyNum = new Random().nextInt(100);
            StringBuilder sb = new StringBuilder(data.length() * copyNum);
            for (int j = 0; j < copyNum; j++) {
                sb.append(data);
            }
            //获取对应数量的令牌
            boolean tokens = tokenBucket.getTokens(sb.toString().getBytes());
            if (tokens) {
                //拿到令牌，通过
                System.out.println("第 " + i + " 次请求通过，传输数据:" + copyNum + "字节 桶内令牌数量：" + tokenBucket.tokenQueue.size());
            } else {
                //没有拿到令牌，拒绝
                System.out.println("第 " + i + " 次请求拒绝，传输数据:" + copyNum + "字节 桶内令牌数量：" + tokenBucket.tokenQueue.size());
            }
            TimeUnit.MILLISECONDS.sleep(100);
        }
    }
}
